package com.example.chatrecommendation.service.impl;

import com.example.chatrecommendation.service.AiPaintingService;
import com.example.chatrecommendation.tool.Api;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.HttpUrl;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.client.methods.RequestBuilder;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.json.JSONArray;
import org.json.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.MalformedURLException;
import java.net.URI;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.ThreadLocalRandom;


@Service
public class AiPaintingServiceImpl implements AiPaintingService {

    private static final Logger log = LoggerFactory.getLogger(AiPaintingServiceImpl.class);
    private static RestTemplate restTemplate = new RestTemplate();
    private final Api api;

    @Autowired
    public AiPaintingServiceImpl(RestTemplate restTemplate, Api api) {
        this.restTemplate = restTemplate;
        this.api = api;
    }


    /** 从响应体中提取Base64字符串
     *
     * @param responseBody
     * @return
     */
    private static String extractBase64FromResponse(String responseBody,String key) {
        ObjectMapper objectMapper = new ObjectMapper();
        try {
            JsonNode jsonNode = objectMapper.readTree(responseBody);
            // 对于spark接口，我们需要从payload.choices.text数组中提取
            if ("spark".equals(key)) {
                JsonNode textArray = jsonNode.path("payload").path("choices").path("text");
                if (textArray.isArray() && textArray.size() > 0) {
                    JsonNode contentNode = textArray.get(0).path("content");
                    if (contentNode.isTextual()) {
                        return contentNode.asText();
                    }
                }
            }
            // 对于sdApi接口，我们直接从images数组中提取
            else if ("sdapi".equals(key)) {
                JsonNode rootNode = objectMapper.readTree(responseBody);
                // 从 JsonNode 对象中提取 "images" 数组
                JsonNode imagesNode = rootNode.get("images");

                // 验证 "images" 是否存在并且是一个数组
                if (imagesNode != null && imagesNode.isArray() && !imagesNode.isEmpty()) {
                    // 提取 "images" 数组中的第一个元素
                    JsonNode firstImageNode = imagesNode.get(0);

                    // 如果存在且为字符串类型，则返回该字符串
                    if (firstImageNode != null && firstImageNode.isTextual()) {
                        return firstImageNode.asText();
                    }
                }
            }
        } catch (Exception e) {
            // 处理异常，例如JSON解析错误
            e.printStackTrace();
        }
        return null; // 如果没有找到或解析失败，返回null
    }

    /**
     * 生成徽章
     * @return
     */
    public String generateHonorBadge() {
        log.info("SD开始生成徽章");

        // 构建请求体
        // seed 1951701925 2287461704 效果最好
        String[] seeds = {"1951701925", "2287461705", "4127465782","3647933624", "3440542522", "3440542523", "3440542525", "1832680489", "3492541044"};
        String randomSeed = seeds[ThreadLocalRandom.current().nextInt(0, seeds.length)];
        log.info("使用的徽章种子：{}", randomSeed);
        String requestJson = """
                {
                  "prompt": "<lora:badge:1>(masterpiece),(best quality),(ultra-detailed),badge,An imaginative and prestigious achievement badge for completing daily tasks consecutively for30 days,featuring a phoenix risirsing from ashes,a calendar,anda trophy.The badge has an ethereal glowing effect with a silver sheen, high-resolution details,and a color scheme of red and gold,symbolizing perseverance and triumph.It's designed with a mythical and victorious style, representing the transfor mative journey of consistent dedication an dsuccess. ",
                  "negative_prompt": "(painting by bad-artist-anime:0.9), (painting by bad-artist:0.9), watermark, text, error, blurry, jpeg artifacts, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, artist name, (worst quality, low quality:1.4), bad anatomy",
                  "cfg_scale": 7,
                  "height": 256,
                  "width": 256,
                  "steps": 30,
                  "sampler_index": "DPM++ 2M Karras",
                  "seed": """ + randomSeed +
                """
                        }
                        """;

        return generateImgs(requestJson);
    }

    /**
     * 生成图像
     * @param prompt
     * @return
     */
    public String generateImage(String prompt) throws MalformedURLException {
        log.info("SD开始生成图像");
        //预处理提示词，删除prompt中开头的
        String preprocessedPrompt = prompt.replaceAll("\\s+", " ");

        if (preprocessedPrompt.contains("<lora:")) {
            return generateHonorBadge();
        }

        prompt = api.translate(preprocessedPrompt);

        // 构建请求体
        String requestJson = """
                {
                  "prompt":"(best quality,4k,8k,highres,masterpiece:1.2),ultra-detailed,(realistic,photorealistic,photo-realistic:1.37),""" + prompt + """
                ",
                  "negative_prompt": "(painting by bad-artist-anime:0.9), (painting by bad-artist:0.9), watermark, text, error, blurry, jpeg artifacts, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, artist name, (worst quality, low quality:1.4), bad anatomy",
                  "cfg_scale": 7.5,
                  "steps": 20,
                  "sampler_index": "DPM++ 2M Karras",
                  "restore_faces": true,
                  "seed": -1
                }
                """;
        log.debug("请求体：\n{}", requestJson);

        return generateImgs(requestJson);
    }

    /**
     * 向ai绘画后端发送请求
     * @param requestJson
     * @return
     */
    private String generateImgs(String requestJson) {
        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);

        // 创建HttpEntity
        HttpEntity<String> entity = new HttpEntity<>(requestJson, headers);

        ResponseEntity<String> response = restTemplate.postForEntity("http://127.0.0.1:7860/sdapi/v1/txt2img", entity, String.class);

        // 这里应该是从response.getBody()中提取Base64字符串
        String imageBase64 = extractBase64FromResponse(response.getBody(), "sdapi");

        return imageBase64;
    }
}
